﻿#include <algorithm>
#include <gtest/gtest.h>
#include <iostream>
#include <vector>

#include <learn/math/differential.h>
#include <learn/math/line_search.h>
#include <learn/math/newton.h>

using namespace xi;
using namespace xi::math;

namespace
{

VectorXd rosenbrock_x0({-10, -10}); // 初始位置
VectorXd rosenbrock_d({1, 1});      // 方向
VectorXd beale_x0({-3, -8.5});      // 初始位置
VectorXd beale_d({2, 3});           // 方向

double parabola(VectorXd x)
{
    return x[0] * x[0] + x[1] * x[1];
}

VectorXd derivative_parabola(VectorXd x)
{
    return {2 * x[0], 2 * x[1]};
}

MatrixXd hessian_parabola(VectorXd x)
{
    MatrixXd H(2, 2);
    H(0, 0) = 2;
    H(0, 1) = 0;
    H(1, 0) = 0;
    H(1, 1) = 2;
    return H;
}

// 最小值点 (1, 1)
double rosenbrock(VectorXd x)
{
    return (1 - x[0]) * (1 - x[0]) + 100 * pow(x[1] - x[0] * x[0], 2);
}

VectorXd derivative_rosenbrock(VectorXd x)
{
    return {-2 * (1 - x[0]) - 400 * x[0] * (x[1] - x[0] * x[0]), 200 * (x[1] - x[0] * x[0])};
}

MatrixXd hessian_rosenbrock(VectorXd x)
{
    MatrixXd H(2, 2);
    H(0, 0) = 2 - 400 * x[1] + 1200 * x[0] * x[0];
    H(0, 1) = -400 * x[0];
    H(1, 0) = -400 * x[0];
    H(1, 1) = 200;
    return H;
}

// 最小值点 (3, 0.5)
double beale(VectorXd x)
{
    return pow(1.5 - x[0] + x[0] * x[1], 2) + pow(2.25 - x[0] + x[0] * x[1] * x[1], 2) +
           pow(2.625 - x[0] + x[0] * x[1] * x[1] * x[1], 2);
}

VectorXd derivative_beale(VectorXd x)
{
    return {2 * (1.5 - x[0] + x[0] * x[1]) * (x[1] - 1) + 2 * (2.25 - x[0] + x[0] * x[1] * x[1]) * (x[1] * x[1] - 1) +
                2 * (2.625 - x[0] + x[0] * x[1] * x[1] * x[1]) * (x[1] * x[1] * x[1] - 1),
            2 * (1.5 - x[0] + x[0] * x[1]) * x[0] + 2 * (2.25 - x[0] + x[0] * x[1] * x[1]) * x[0] * x[1] * 2 +
                2 * (2.625 - x[0] + x[0] * x[1] * x[1] * x[1]) * x[0] * x[1] * x[1] * 3};
}

} // namespace

// 测试 Jacobian 函数
TEST(TestDifferential, diff)
{
    auto f = [](VectorXd x) -> double { return 2 * sin(x[0]) + cos(x[1]); };
    auto g = [](VectorXd x) -> VectorXd { return {2 * sin(x[0]) + cos(x[1]), -2 * cos(x[0]) + sin(x[1])}; };

    auto fj = jacobian<double>(f, {0, 0});
    auto gj = jacobian<VectorXd>(g, {0, 0});

    MatrixXd fJ(1, 2);
    fJ(0, 0) = 2;
    fJ(0, 1) = 0;
    EXPECT_LT((fJ - fj).norm(), Numeric::distance);

    MatrixXd gJ(2, 2);
    gJ(0, 0) = 2;
    gJ(0, 1) = 0;
    gJ(1, 0) = 0;
    gJ(1, 1) = 1;
    // std::cout << gJ << std::endl;
    // std::cout << gj << std::endl;
    EXPECT_LT((gJ - gj).norm(), Numeric::distance);
}

TEST(TestDifferential, hessen)
{
    auto f = [](VectorXd x) -> double { return 10 * sin(x[0]) + 3 * cos(x[1]); };
    auto h = hessian(f, {0, 0});
    std::cout << h << std::endl;
}

TEST(TestLineSearch, GoldenSection)
{
    {
        double x1 = golden_section_line_search(parabola, VectorXd{-1, -1}, VectorXd{1, 1}, 1e-8);
        double x2 = golden_section_line_search(parabola, VectorXd{1, 1}, VectorXd{1, 1}, 1e-8);

        EXPECT_LE(std::abs(x1 - 1), Numeric::distance);
        EXPECT_LE(std::abs(x2 + 1), Numeric::distance);
    }

    {
        auto x1 = golden_section_line_search(rosenbrock, rosenbrock_x0, rosenbrock_d, 1e-8);
        auto x2 = golden_section_line_search(rosenbrock, rosenbrock_x0, VectorXd(-rosenbrock_d), 1e-8);

        EXPECT_LE(std::abs(x1 - 11), Numeric::distance);
        EXPECT_LE(std::abs(x2 + 11), Numeric::distance);
    }

    {
        double x1 = golden_section_line_search(beale, beale_x0, beale_d, 1e-8);
        double x2 = golden_section_line_search(beale, beale_x0, VectorXd(-beale_d), 1e-8);

        EXPECT_LE(std::abs(x1 - 3), Numeric::distance);
        EXPECT_LE(std::abs(x2 + 3), Numeric::distance);
    }
}

TEST(TestLineSearch, Armijo)
{
    {
        double x1 = armijo_goldstein_line_search(parabola, derivative_parabola, VectorXd{-1, -1}, VectorXd{1, 1}, 1e-8);
        // double x2 = armijo_goldstein_line_search(parabola, derivative_parabola, VectorXd{1, 1}, VectorXd{1, 1});

        EXPECT_LE(std::abs(x1 - 1), 1);
        // EXPECT_LE(std::abs(x2 + 1), 1);
    }

    {
        double x1 = armijo_goldstein_line_search(rosenbrock, derivative_rosenbrock, rosenbrock_x0, rosenbrock_d, 1e-8);
        // double x2 = armijo_goldstein_line_search(rosenbrock, derivative_rosenbrock, rosenbrock_x0,
        // VectorXd(-rosenbrock_d));

        EXPECT_LE(std::abs(x1 - 11), 1);
        // EXPECT_LE(std::abs(x2 + 11), 1);
    }

    {
        double x1 = armijo_goldstein_line_search(beale, derivative_beale, beale_x0, beale_d, 1e-8);
        // double x2 = armijo_goldstein_line_search(beale, derivative_beale, beale_x0, VectorXd(-beale_d));

        EXPECT_LE(std::abs(x1 - 3), 3);
        // EXPECT_LE(std::abs(x2 + 3), 1);
    }
}

TEST(TestLineSearch, Wolfe)
{
    {
        double x1 = wolfe_powell_line_search(parabola, derivative_parabola, VectorXd{-1, -1}, VectorXd{1, 1}, 1e-8);
        // double x2 = wolfe_powell_line_search(parabola, derivative_parabola, VectorXd{1, 1}, VectorXd{1, 1});

        EXPECT_LE(std::abs(x1 - 1), 1);
        // EXPECT_LE(std::abs(x2 + 1), 1);
    }

    {
        double x1 = wolfe_powell_line_search(rosenbrock, derivative_rosenbrock, rosenbrock_x0, rosenbrock_d, 1e-8);
        // double x2 = wolfe_powell_line_search(rosenbrock, derivative_rosenbrock, rosenbrock_x0,
        // VectorXd(-rosenbrock_d));

        EXPECT_LE(std::abs(x1 - 11), 1);
        // EXPECT_LE(std::abs(x2 + 11), 1);
    }

    {
        double x1 = wolfe_powell_line_search(beale, derivative_beale, beale_x0, beale_d, 1e-8);
        // double x2 = wolfe_powell_line_search(beale, derivative_beale, beale_x0, VectorXd(-beale_d));

        EXPECT_LE(std::abs(x1 - 3), 3);
        // EXPECT_LE(std::abs(x2 + 3), 1);
    }
}

TEST(TestLineSearch, Backtracking)
{
    {
        double x1 = backtracking_line_search(parabola, derivative_parabola, VectorXd{-1, -1}, VectorXd{1, 1}, 1e-8);
        // double x2 = backtracking_line_search(parabola, derivative_parabola, VectorXd{1, 1}, VectorXd{1, 1});

        EXPECT_LE(std::abs(x1 - 1), 1);
        // EXPECT_LE(std::abs(x2 + 1), 1);
    }

    {
        double x1 = backtracking_line_search(rosenbrock, derivative_rosenbrock, rosenbrock_x0, rosenbrock_d, 1e-8);
        // double x2 = backtracking_line_search(rosenbrock, derivative_rosenbrock, rosenbrock_x0,
        // VectorXd(-rosenbrock_d));

        EXPECT_LE(std::abs(x1 - 11), 1);
        // EXPECT_LE(std::abs(x2 + 11), 1);
    }

    {
        double x1 = backtracking_line_search(beale, derivative_beale, beale_x0, beale_d, 1e-8);
        // double x2 = backtracking_line_search(beale, derivative_beale, beale_x0, VectorXd(-beale_d));

        EXPECT_LE(std::abs(x1 - 3), 3);
        // EXPECT_LE(std::abs(x2 + 3), 1);
    }
}

TEST(TestLineSearch, Steepest)
{
    auto x1 = steepest_descent(parabola, derivative_parabola, VectorXd{-1, -1}, 1e-10);
    EXPECT_LE(x1.norm(), Numeric::distance);

    auto x3 = steepest_descent(rosenbrock, derivative_rosenbrock, rosenbrock_x0, 1e-10);
    auto x4 = steepest_descent(beale, derivative_beale, beale_x0, 1e-10);
    EXPECT_LE((x3 - VectorXd{1, 1}).norm(), Numeric::distance);
    EXPECT_LE((x4 - VectorXd{3, 0.5}).norm(), Numeric::distance);
}

TEST(TestLineSearch, Newton)
{
    auto x1 = pure_newton(derivative_parabola, hessian_parabola, VectorXd{-1, -1}, 1e-10);
    EXPECT_LE(x1.norm(), Numeric::distance);

    auto x2 = pure_newton(derivative_rosenbrock, hessian_rosenbrock, rosenbrock_x0, 1e-10);
    EXPECT_LE((x2 - VectorXd{1, 1}).norm(), Numeric::distance);
}

TEST(TestLineSearch, LineNewton)
{
    auto x1 = line_search_newton(parabola, derivative_parabola, hessian_parabola, VectorXd{-1, -1}, 1e-10);
    EXPECT_LE(x1.norm(), Numeric::distance);

    auto x2 = line_search_newton(rosenbrock, derivative_rosenbrock, hessian_rosenbrock, rosenbrock_x0, 1e-10);
    EXPECT_LE((x2 - VectorXd{1, 1}).norm(), Numeric::distance);
}